﻿#include <osg/Image>
#include <osg/Notify>
#include <osg/Geode>
#include <osg/GL>

#include <osgDB/Registry>
#include <osgDB/FileNameUtils>
#include <osgDB/FileUtils>
#include <osgDB/ImageOptions>

#include <osgTerrain/TerrainTile>
#include <osgTerrain/Terrain>
#include <osg/PagedLOD>

#include <osgDB/ReadFile>
#include <osgDB/WriteFile>
#include <osg/Texture2D>
#include <osg/PagedLOD>
#include <queue>

#include <stdlib.h>
#include "Attributes.h"
#include "json.hpp"

enum TYPE
{
    NORMAL = 0,
    LEAF = 1,
    PROXY = 2,
};

class HNode;
class HNode : public osg::Referenced
{
public:
    HNode::HNode()
    {
        _name = "";
        _numPoints = 0;
        _childMask = 0;
        _byteOffset = 0;
        _byteSize = 0;
        _type = TYPE::LEAF;
        _children.resize(8);
        for (int i = 0; i < 8; i++)
        {
            _children[i] = NULL;
        }
    }

    std::string   _name = "";
    int      _numPoints = 0;
    uint8_t  _childMask = 0;
    uint64_t _byteOffset = 0;
    uint64_t _byteSize = 0;
    TYPE     _type = TYPE::LEAF;
    std::vector<osg::ref_ptr<HNode> > _children;
};

class Container : public osg::Object
{
public:
    Container()
    {
    }

    Container(const Container& other, const osg::CopyOp& copyop = osg::CopyOp::SHALLOW_COPY) : osg::Object(other, copyop)
    {
        _hNode = other._hNode;
    }

    explicit Container(osg::ref_ptr<HNode> hNode, osg::ref_ptr<Attributes> attributes)
    {
        _hNode = hNode;
        _attributes = attributes;
    }

    osg::ref_ptr<HNode> node() const
    {
        return _hNode;
    }

    osg::ref_ptr<Attributes> attributes()  const
    {
        return _attributes;
    }

    META_Object(OSG, Container)

protected:
    osg::ref_ptr<HNode> _hNode;
    osg::ref_ptr<Attributes> _attributes;
};

class ReaderWriterPotree : public osgDB::ReaderWriter
{
public:
    ReaderWriterPotree()
    {
        supportsExtension("potree", "potree reader");
        supportsExtension("pchildren", "Internal use of potree <children> tag");
    }

    virtual const char* className() const
    {
        return "Potree Reader";
    }

    void ParseToHNode(const std::vector<char>& data, osg::ref_ptr<HNode> node) const
    {
        node->_type = TYPE(data[0]);
        node->_childMask = uint8_t(data[1]);
        node->_numPoints = ((uint32_t*)(data.data() + 2))[0];
        node->_byteOffset = ((uint64_t*)(data.data() + 6))[0];
        node->_byteSize = ((uint64_t*)(data.data() + 14))[0];
    }

    bool ParseHierarchy(const std::string& hierarchyPath, int offset, osg::ref_ptr<HNode>& hNode) const
    {
        std::fstream fin(hierarchyPath, std::ios::binary | std::ios::in);
        if (!fin)
        {
            OSG_WARN << "无法打开文件:" << hierarchyPath << std::endl;
            return false;
        }
        fin.seekg(offset, std::ios::beg);
        hNode = new HNode;
        std::vector<char> data(22, 0);
        fin.read(data.data(), 22);
        ParseToHNode(data, hNode);
        std::list<osg::ref_ptr<HNode> > queues;
        queues.push_back(hNode);
        while (queues.size() > 0)
        {
            osg::ref_ptr<HNode> hNode = queues.front();
            queues.pop_front();
            if (hNode->_type == TYPE::PROXY)
            {
                continue;
            }
            for (int i = 0; i < 8; i++)
            {
                if (hNode->_childMask & (1 << i))
                {
                    hNode->_children[i] = new HNode;
                    std::vector<char> data(22, 0);
                    fin.read(data.data(), 22);
                    ParseToHNode(data, hNode->_children[i]);
                    queues.push_back(hNode->_children[i]);
                }
            }
        }
        fin.close();
        return true;
    }

    bool ParseMetadata(const std::string& path, osg::ref_ptr<Attributes>& attributes) const
    {
        attributes = new Attributes;
        std::ifstream file(path);
        if (!file.is_open())
        {
            OSG_WARN << "无法打开文件:" << path << std::endl;
            return false;
        }
        nlohmann::json data = nlohmann::json::parse(file);
        file.close();

        //解析BoundingBox
        auto& bbox = data["boundingBox"];
        std::vector<double> bbox_min = bbox["min"].get<std::vector<double>>();
        std::vector<double> bbox_max = bbox["max"].get<std::vector<double>>();
        attributes->_box = osg::BoundingBoxd(osg::Vec3d(bbox_min[0], bbox_min[1], bbox_min[2]), osg::Vec3d(bbox_max[0], bbox_max[1], bbox_max[2]));
        //解析offset, scale
        std::vector<double> offset = data["offset"].get<std::vector<double>>();
        std::vector<double> scale = data["scale"].get<std::vector<double>>();
        attributes->_posOffset = osg::Vec3d(offset[0], offset[1], offset[2]);
        attributes->_posScale = osg::Vec3d(scale[0], scale[1], scale[2]);
        //解析attributes数组
        std::vector<Attribute> vecAttribute;
        for (auto& attr : data["attributes"])
        {
            Attribute a;
            a.name = attr["name"];
            a.description = attr["description"];
            a.size = attr["size"];
            a.numElements = attr["numElements"];
            a.elementSize = attr["elementSize"];
            a.type = typenameToType(attr["type"]);
            std::vector<double> mindata = attr["min"].get<std::vector<double>>();
			if (mindata.size() == 3)
				a.min = osg::Vec3d(mindata[0], mindata[1], mindata[2]);
			else
				a.min = osg::Vec3d(mindata[0], mindata[0], mindata[0]);
            std::vector<double> maxdata = attr["max"].get<std::vector<double>>();
			if (maxdata.size() == 3)
            a.max = osg::Vec3d(maxdata[0], maxdata[1], maxdata[2]);
			else
				a.max = osg::Vec3d(maxdata[0], maxdata[0], maxdata[0]);
            std::vector<double> scaledata = attr["scale"].get<std::vector<double>>();
			if (scaledata.size() == 3)
            a.scale = osg::Vec3d(scaledata[0], scaledata[1], scaledata[2]);
			else
				a.scale = osg::Vec3d(scaledata[0], scaledata[0], scaledata[0]);
            std::vector<double> offsetdata = attr["offset"].get<std::vector<double>>();
			if (offsetdata.size() == 3)
            a.offset = osg::Vec3d(offsetdata[0], offsetdata[1], offsetdata[2]);
			else
				a.offset = osg::Vec3d(offsetdata[0], offsetdata[0], offsetdata[0]);

            // 特殊处理histogram (只有classification有)
            if (attr.contains("histogram"))
            {
                a.histogram = attr["histogram"].get<std::vector<int64_t>>();
            }
            vecAttribute.push_back(a);
        }
        attributes->setAttribute(vecAttribute);
        return true;
    }

    virtual ReadResult readNode(const std::string& location, const Options* options) const
    {
        std::string ext = osgDB::getFileExtension(location);
        if (!acceptsExtension(ext))
            return ReadResult::FILE_NOT_HANDLED;
        std::string filePathDir = osgDB::getNameLessExtension(location);
        if (ext == "pchildren" && options && options->getUserDataContainer())
        {
            osg::ref_ptr<const Container> parentHNodeContainer = dynamic_cast<const Container*> (options->getUserDataContainer()->getUserObject("ParentTile"));
            if (parentHNodeContainer)
            {
                return createTileChildren(parentHNodeContainer->node(), parentHNodeContainer->attributes(), osgDB::getNameLessExtension(filePathDir), options);
            }
        }
        else
        {
            std::string hierarchyPath = osgDB::concatPaths(filePathDir, "hierarchy.bin");
            osg::ref_ptr<HNode> hNode;
            if (ParseHierarchy(hierarchyPath, 0, hNode))
            {
                osg::ref_ptr<Attributes> attributes = NULL;
                ParseMetadata(osgDB::concatPaths(filePathDir, "metadata.json"), attributes);
                return createTile(hNode, attributes, filePathDir, options);
            }
        }
        return NULL;
    }

    osg::Node* createTileChildren(osg::ref_ptr<HNode> parentHNode, osg::ref_ptr<Attributes> attributes, const std::string& filePathDir, const Options* options) const
    {
        if (!parentHNode)
        {
            return NULL;
        }

        osg::ref_ptr<osg::Group> grp = new osg::Group;
        for (int i = 0; i < 8; i++)
        {
            if (parentHNode->_children[i])
            {
                osg::ref_ptr<osg::Node> node = createTile(parentHNode->_children[i], attributes, filePathDir, options);
                if (node)
                {
                    grp->addChild(node);
                }
            }
        }
        return grp.release();
    }

    osg::Node* createPointCloudNode(const std::vector<char>& data, int pointNum, osg::ref_ptr<Attributes> attributes) const
    {
        osg::Vec3d center = attributes->_box.center();
        osg::ref_ptr<osg::MatrixTransform> mt = new osg::MatrixTransform(osg::Matrix::translate(center));
        osg::ref_ptr<osg::Geometry> geometry = new osg::Geometry;
        mt->addChild(geometry);
        osg::ref_ptr<osg::Vec3Array> v3a = new osg::Vec3Array;
        osg::ref_ptr<osg::Vec4ubArray> v4c = new osg::Vec4ubArray;
        for (int i = 0; i < pointNum; i++)
        {
            const char* pData = data.data() + (attributes->_bytes * i);
            for (int j = 0; j < attributes->_list.size(); j++)
            {
                if ("position" == attributes->_list[j].name)
                {
                    if (INT32 == attributes->_list[j].type || UINT32 == attributes->_list[j].type)
                    {
                        int* pDataInt = (int*)pData;
                        osg::Vec3d position = osg::Vec3d(attributes->_posOffset.x() + attributes->_posScale.x() * pDataInt[0] - center.x(),
                            attributes->_posOffset.y() + attributes->_posScale.y() * pDataInt[1] - center.y(),
                            attributes->_posOffset.z() + attributes->_posScale.z() * pDataInt[2] - center.z());
                        v3a->push_back(position);
                    }
                    else
                    {
                        OSG_FATAL << "position type is:" << attributes->_list[j].type << ", please improve" << std::endl;
                    }
                }
                else if ("rgb" == attributes->_list[j].name)
                {
                    if (INT16 == attributes->_list[j].type || UINT16 == attributes->_list[j].type)
                    {
                        uint16_t* pDataUInt16 = (uint16_t*)pData;
                        osg::Vec4ub color = osg::Vec4ub(pDataUInt16[0] / 255,
                            pDataUInt16[1] / 255,
                            pDataUInt16[2] / 255,
                            255);
                        v4c->push_back(color);
                    }
                    else
                    {
                        OSG_FATAL << "color type is:" << attributes->_list[j].type << ", please improve" << std::endl;
                    }
                }
                pData += attributes->_list[j].size;
            }
        }
        geometry->setVertexArray(v3a);
        if (v4c->size() > 0)
        {
            geometry->setColorArray(v4c, osg::Array::BIND_PER_VERTEX);
        }
        geometry->addPrimitiveSet(new osg::DrawArrays(osg::PrimitiveSet::POINTS, 0, v3a->size()));
        return mt.release();
    }

    osg::Node* createTile(osg::ref_ptr<HNode> hNode, osg::ref_ptr<Attributes> attributes, const std::string& filePathDir, const Options* options) const
    {
        if (!hNode)
        {
            return NULL;
        }
        std::string binPath = osgDB::concatPaths(filePathDir, "octree.bin");
        if (TYPE::PROXY == hNode->_type)
        {
            std::string hierarchyPath = osgDB::concatPaths(filePathDir, "hierarchy.bin");
            osg::ref_ptr<HNode> hChunkNode = NULL;
            if (ParseHierarchy(hierarchyPath, hNode->_byteOffset, hChunkNode))
            {
                return createTile(hChunkNode, attributes, filePathDir, options);
            }
        }
        else
        {
            std::fstream fin(binPath, std::ios::binary | std::ios::in);
            if (!fin)
            {
                OSG_WARN << "无法打开文件:" << binPath << std::endl;
                return NULL;
            }
            fin.seekg(hNode->_byteOffset, std::ios::beg);
            std::vector<char> data(hNode->_byteSize, 0);
            fin.read(data.data(), hNode->_byteSize);
            fin.close();
            osg::ref_ptr<osg::Node> pointcloudNode = createPointCloudNode(data, hNode->_numPoints, attributes);
            if (0 != hNode->_childMask)
            {
                osg::ref_ptr<osg::PagedLOD> plod = new osg::PagedLOD;
                plod->addChild(pointcloudNode);
                osgDB::Options* childOpt = new osgDB::Options;
                osg::ref_ptr<Container> parentContainer = new Container(hNode, attributes);
                parentContainer->setName("ParentTile");
                childOpt->getOrCreateUserDataContainer()->addUserObject(parentContainer);
                plod->setDatabaseOptions(childOpt);
                plod->setFileName(1, filePathDir + "." + std::to_string(hNode->_byteOffset) + ".pchildren");

                if (pointcloudNode.valid())
                {
                    osg::BoundingSphered bound2;// = bound;  // FIXME: some <boundingVolume> too far away?
                    osg::BoundingSphere bound0 = pointcloudNode->getBound();
                    bound2.expandBy(osg::BoundingSphered(bound0.center(), bound0.radius()));
                    plod->setCenterMode(osg::LOD::USER_DEFINED_CENTER);
                    plod->setCenter(bound2.center()); 
                    plod->setRadius(bound2.radius());
                }
                else
                {
                    OSG_WARN << "[ReaderWriterPotree] Missing <boundingVolume>?" << std::endl;
                }
                //[Modify by author for Issue https://gitee.com/osg_opensource/osg-potree/issues/ICL4D1 at 20250712 
                //plod->setRangeMode(osg::LOD::PIXEL_SIZE_ON_SCREEN);
                //plod->setRange(0, 0.0f, FLT_MAX);
                //plod->setRange(1, 200, FLT_MAX);
                plod->setRangeMode(osg::LOD::DISTANCE_FROM_EYE_POINT);
                plod->setRange(0, 0.0f, FLT_MAX);
                double radius = plod->getRadius();
                double range = 1080 * std::sqrt(4 * radius * radius / hNode->_numPoints) /2/ 0.5629;//假设屏幕高度像素为1080//数据在屏幕的最大截面面积为4 * radius * radius//视场角为60，tan(30)=0.5629
                plod->setRange(1, 0, range);
                //]Modify by author for Issue https://gitee.com/osg_opensource/osg-potree/issues/ICL4D1 at 20250712 
                return plod.release();
            }
            else
            {
                return pointcloudNode.release();
            }
        }

        return NULL;
    }

    virtual WriteResult writeNode(const osg::Node& node, const std::string& location, const Options* options) const
    {
        std::string ext = osgDB::getFileExtension(location);
        if (!acceptsExtension(ext))
            return WriteResult::FILE_NOT_HANDLED;
        OSG_WARN << "not implemented" << std::endl;
        return WriteResult::FILE_SAVED;
    }

};

REGISTER_OSGPLUGIN(ifc, ReaderWriterPotree)
